import akshare as ak
import pandas as pd
from sqlalchemy import create_engine, text
import time
import logging
from tqdm import tqdm

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('stock_info_xq.log'),
        logging.StreamHandler()
    ]
)

# 数据库配置
DATABASE_URL = "postgresql://leizhen:751982leizhen@localhost/stock_db"
engine = create_engine(DATABASE_URL)

def get_stock_symbols():
    """从数据库获取股票代码列表"""
    try:
        df = pd.read_sql("SELECT symbol FROM stocks", engine)
        return df['symbol'].tolist()
    except Exception as e:
        logging.error(f"获取股票列表失败: {str(e)}")
        return []

def add_prefix(symbol):
    """为股票代码添加交易所前缀"""
    if symbol.startswith('6'):
        return f"SH{symbol}"
    elif symbol.startswith(('0', '3')):
        return f"SZ{symbol}"
    else:
        logging.warning(f"未知股票代码格式: {symbol}")
        return None

def get_stock_info(symbol, max_retries=3):
    """获取单个股票详细信息"""
    prefixed_symbol = add_prefix(symbol)
    print(f"处理股票代码: {prefixed_symbol}")
    if not prefixed_symbol:
        return None
        
    for attempt in range(max_retries):
        try:
            # 使用akshare获取数据并处理不同格式
            data = ak.stock_individual_basic_info_xq(symbol=prefixed_symbol)
            # 处理不同返回格式
            if isinstance(data, pd.DataFrame):
                df = data.copy()
            elif isinstance(data, dict):
                if "data" in data:
                    df = pd.DataFrame(data["data"])
                else:
                    df = pd.DataFrame([data])
            else:
                raise ValueError("未知的数据格式")
            
            if df.empty:
                raise ValueError("返回数据为空")
                
            # 确保数据格式正确
            required_columns = ['item', 'value']
            if not all(col in df.columns for col in required_columns):
                raise ValueError(f"数据缺少必要列: {required_columns}")
            
            df['symbol'] = symbol  # 保留原始股票代码
            
            # 打印将要保存的数据
            logging.info(f"准备保存{symbol}的数据, 共{len(df)}条")
            
            return df
                
        except Exception as e:
            logging.warning(f"第{attempt+1}次尝试获取{symbol}信息失败: {str(e)}")
            if attempt < max_retries - 1:
                time.sleep(2)  # 等待2秒后重试
                continue
            logging.error(f"获取{symbol}信息最终失败: {str(e)}", exc_info=True)
            return None

def init_table():
    """初始化股票详细信息表"""
    max_retries = 3
    for attempt in range(max_retries):
        try:
            with engine.connect() as conn:
                # 检查表是否存在
                table_exists = conn.execute(text("""
                    SELECT EXISTS (
                        SELECT FROM information_schema.tables 
                        WHERE table_name = 'stock_individual_basic_info_xq'
                    )
                """)).scalar()
                
                if not table_exists:
                    conn.execute(text("""
                        CREATE TABLE stock_individual_basic_info_xq (
                            symbol VARCHAR(10) PRIMARY KEY,
                            org_id VARCHAR(20),
                            org_name_cn TEXT,
                            org_short_name_cn TEXT,
                            org_name_en TEXT,
                            org_short_name_en TEXT,
                            main_operation_business TEXT,
                            operating_scope TEXT,
                            district_encode VARCHAR(20),
                            org_cn_introduction TEXT,
                            legal_representative TEXT,
                            general_manager TEXT,
                            secretary TEXT,
                            established_date BIGINT,
                            reg_asset FLOAT,
                            staff_num INTEGER,
                            telephone TEXT,
                            postcode TEXT,
                            fax TEXT,
                            email TEXT,
                            org_website TEXT,
                            reg_address_cn TEXT,
                            reg_address_en TEXT,
                            office_address_cn TEXT,
                            office_address_en TEXT,
                            currency_encode VARCHAR(20),
                            currency TEXT,
                            listed_date BIGINT,
                            provincial_name TEXT,
                            actual_controller TEXT,
                            classi_name TEXT,
                            pre_name_cn TEXT,
                            chairman TEXT,
                            executives_nums INTEGER,
                            actual_issue_vol FLOAT,
                            issue_price FLOAT,
                            actual_rc_net_amt FLOAT,
                            pe_after_issuing FLOAT,
                            online_success_rate_of_issue FLOAT,
                            affiliate_industry TEXT,
                            update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                        )
                    """))
                    conn.commit()
                    logging.info("股票详细信息表创建成功")
                else:
                    logging.info("股票详细信息表已存在")
                return
        except Exception as e:
            if attempt == max_retries - 1:
                logging.error(f"初始化表失败: {str(e)}")
                raise
            time.sleep(1)

def save_stock_info(symbol, info_df):
    """保存股票完整信息"""
    max_retries = 3
    for attempt in range(max_retries):
        try:
            with engine.begin() as conn:
                # 构建字段映射
                field_map = {
                    'org_id': 'org_id',
                    'org_name_cn': 'org_name_cn',
                    'org_short_name_cn': 'org_short_name_cn',
                    'org_name_en': 'org_name_en',
                    'org_short_name_en': 'org_short_name_en',
                    'main_operation_business': 'main_operation_business',
                    'operating_scope': 'operating_scope',
                    'district_encode': 'district_encode',
                    'org_cn_introduction': 'org_cn_introduction',
                    'legal_representative': 'legal_representative',
                    'general_manager': 'general_manager',
                    'secretary': 'secretary',
                    'established_date': 'established_date',
                    'reg_asset': 'reg_asset',
                    'staff_num': 'staff_num',
                    'telephone': 'telephone',
                    'postcode': 'postcode',
                    'fax': 'fax',
                    'email': 'email',
                    'org_website': 'org_website',
                    'reg_address_cn': 'reg_address_cn',
                    'reg_address_en': 'reg_address_en',
                    'office_address_cn': 'office_address_cn',
                    'office_address_en': 'office_address_en',
                    'currency_encode': 'currency_encode',
                    'currency': 'currency',
                    'listed_date': 'listed_date',
                    'provincial_name': 'provincial_name',
                    'actual_controller': 'actual_controller',
                    'classi_name': 'classi_name',
                    'pre_name_cn': 'pre_name_cn',
                    'chairman': 'chairman',
                    'executives_nums': 'executives_nums',
                    'actual_issue_vol': 'actual_issue_vol',
                    'issue_price': 'issue_price',
                    'actual_rc_net_amt': 'actual_rc_net_amt',
                    'pe_after_issuing': 'pe_after_issuing',
                    'online_success_rate_of_issue': 'online_success_rate_of_issue',
                    'affiliate_industry': 'affiliate_industry'
                }
                
                # 构建参数字典
                params = {'symbol': symbol}
                for _, row in info_df.iterrows():
                    item = row['item']
                    value = row['value']
                    if item in field_map:
                        if isinstance(value, dict):
                            import json
                            params[field_map[item]] = json.dumps(value, ensure_ascii=False)
                        else:
                            params[field_map[item]] = str(value) if value is not None else None
                
                # 构建动态SQL
                columns = ['symbol'] + [field for field in field_map.values()]
                values = [f':{col}' for col in columns]
                set_clause = ', '.join([f'{col} = EXCLUDED.{col}' for col in columns if col != 'symbol'])
                
                sql = f"""
                    INSERT INTO stock_individual_basic_info_xq ({', '.join(columns)})
                    VALUES ({', '.join(values)})
                    ON CONFLICT (symbol) DO UPDATE
                    SET {set_clause},
                        update_time = CURRENT_TIMESTAMP
                """
                
                conn.execute(text(sql), params)
                return True
        except Exception as e:
            if attempt == max_retries - 1:
                logging.error(f"保存股票{symbol}信息最终失败: {str(e)}")
                return False
            time.sleep(1)
    return False

def main():
    init_table()
    
    symbols = get_stock_symbols()
    if not symbols:
        return
        
    success_count = 0
    for symbol in tqdm(symbols, desc="获取股票信息"):
        info = get_stock_info(symbol)
        if info is None:
            continue
            
        if save_stock_info(symbol, info):
            logging.info(f"成功保存股票{symbol}的信息")
            success_count += 1
            
        time.sleep(1)  # 控制请求频率
    
    logging.info(f"数据采集完成，共保存{success_count}只股票的信息")

if __name__ == "__main__":
    main()
